from llmweb.base import LLMWebBase, LoginMixin
from llmweb import BrowserPool
from typing import Optional, AsyncGenerator
import asyncio
import time
from llmweb.logging import logger

async def keep_browser_open(page, timeout):
    # 这个函数会等待指定的时间，直到超时或被取消
    await asyncio.sleep(timeout)
class YiYanLLM(LLMWebBase, LoginMixin):
    def __init__(self, pool: Optional[BrowserPool]=BrowserPool()):
        site = "https://yiyan.baidu.com/"
        super().__init__(site, pool) 
    
    async def generate(self, prompt: str, **kwargs) -> AsyncGenerator[str, None]:
        page = await self.pool.acquire_page()
        stop = kwargs.get("stop")
        logger.debug(f"get stop:{stop}")
        await page.goto(self.site)
        try:
            await self.pool.load_local_storage(page)
            await self.pool.load_cookies(page)
            await page.reload()
        except Exception as e:
            logger.error(f"Load local storage error: {e}")
            await self.pool.release_page(page)
            raise e
        
        await page.wait_for_load_state('networkidle')
        await page.get_by_text("文心大模型4.0").last.click()
        element = page.locator("#eb_model_footer").get_by_role("textbox")
        await element.fill(prompt)
        span_elements = page.locator("#eb_model_footer").locator("div >> div >> div >> div >> span")
        await span_elements.last.click()
        await page.wait_for_selector('div.custom-html')
        els = page.locator("div.custom-html")
        text: str = await els.last.inner_text()
        if len(text) > 0:
            yield text
            
        total_timeout_time = time.time() + self.wait_timeout
        chunk_wait_timeout = time.time() + 5 # 5秒
        while True:
            try:
                if time.time() > total_timeout_time:
                    ## 总体超时，退出循环
                    raise TimeoutError("text response timeout")
                new_text: str = await els.last.inner_text()
                if stop is not None and isinstance(stop,list):
                    for s in stop:
                        if s in new_text:
                            await page.get_by_text("停止生成").click()
                            logger.debug("stop text")
                            raise StopIteration(s)
                if text != new_text:
                    chunks = new_text.replace(text, "")
                    yield chunks
                    text = new_text
                    chunk_wait_timeout = time.time() + 5
                else:
                    if time.time() > chunk_wait_timeout:
                        #等待单个块超时，退出循环
                        break
                        
                await asyncio.sleep(0.5)
            except Exception as e:
                logger.error(e)
                break
        await self.pool.release_page(page)
        logger.debug("exit yiyan generator")
        
    async def login(self, **kwargs):
        """
        需要人工登录，验证，存储登录信息
        """
        page = await self.pool.acquire_page()
        await page.goto("https://yiyan.baidu.com/")
        page.on('load', lambda: asyncio.create_task(self.pool.add_save_cookie_btn(page)))
        await asyncio.sleep(40)
        await self.pool.release_page(page)
